//////////////////////////////////////////////////////////////////////////////
//// Original Copyright (c) Microsoft Corporation. All rights reserved
//// This software contains source code provided by NVIDIA Corporation.
//////////////////////////////////////////////////////////////////////////////

//----------------------------------------------------------------------------
// File: MersenneTwister.h
// 
// This sample implements Mersenne Twister random number generator 
// and Cartesian Box-Muller transformation on the GPU.
//----------------------------------------------------------------------------

#include <math.h>
#include <fstream>
#include <iostream>
#include <amp.h>
#include <amp_math.h>
#include <ppl.h>

#pragma warning (disable : 4267)

using namespace concurrency;

class MSParallelMersenneTwister
{
private:
    static int div_up(int a, int b)
    {
        return (a + b - 1)/b;
    }

    //Align a to nearest higher multiple of b
    static int align_up(int a, int b)
    {
        return ((a + b - 1)/b) * b;
    }

    typedef struct
    {
        unsigned int matrix_a;
        unsigned int mask_b;
        unsigned int mask_c;
        unsigned int seed;
    } mt_struct;

    #define   MT_RNG_COUNT 4096
    #define          MT_MM 9
    #define          MT_NN 19
    #define       MT_WMASK 0xFFFFFFFFU
    #define       MT_UMASK 0xFFFFFFFEU
    #define       MT_LMASK 0x1U
    #define      MT_SHIFT0 12
    #define      MT_SHIFTB 7
    #define      MT_SHIFTC 15
    #define      MT_SHIFT1 18

    //Load twister configurations
    void load_MT(const char *fname, mt_struct* data, int elements)
    {
        FILE *fd;
        if(0 != fopen_s(&fd, fname, "rb"))
        {
            printf("load_MT(): failed to open %s\n", fname);
            printf("TEST FAILED\n");
            exit(0);
        }
        if( fread(data, elements * sizeof(mt_struct), 1, fd) == elements * sizeof(mt_struct))
        {
            printf("load_MT(): failed to load %s\n", fname);
            printf("TEST FAILED\n");
            exit(0);
        }
        fclose(fd);
    }

    //Initialize/seed twister
    void seed_MT(unsigned int seed0, mt_struct* data, std::vector<unsigned int>& matrix, 
                 std::vector<unsigned int>& mask_b, std::vector<unsigned int>& mask_c, 
                 std::vector<unsigned int>& seed) 
    {
        for(int i = 0; i < MT_RNG_COUNT; i++)
        {
            matrix[i] = data[i].matrix_a;
            mask_b[i] = data[i].mask_b;
            mask_c[i] = data[i].mask_c;
            seed[i] = seed0;
        }
    }

public:
    ////////////////////////////////////////////////////////////////////////////////
    // Write MT_RNG_COUNT vertical lanes of n_per_RNG random numbers to random.
    // For coalesced global writes MT_RNG_COUNT should be a multiple of hardware scheduling unit size.
    // Hardware scheduling unit is called warp or wave or wavefront
    // Initial states for each generator are the same, since the states are
    // initialized from the global seed. In order to improve distribution properties
    // on small n_per_RNG supply dedicated (local) seed to each twister.
    // The local seeds, in their turn, can be extracted from global seed
    // by means of any simple random number generator, like LCG.
    ////////////////////////////////////////////////////////////////////////////////
    static void rand_MT_kernel(
        int idx,
        array_view<float, 2> random, 
        const unsigned int matrix_a, 
        const unsigned int mask_b, const unsigned int mask_c, 
        const unsigned int seed, const int n_per_RNG
    ) restrict(cpu, amp)
    {
        int state_1;
        int state_M;
        unsigned int mti, mti_M, x;
        unsigned int mti_1, mt[MT_NN];

        //Bit-vector Mersenne Twister parameters are in matrix_a, mask_b, mask_c, seed
        //Initialize current state
        mt[0] = seed;
        for(int state = 1; state < MT_NN; state++)
            mt[state] = (1812433253U * (mt[state - 1] ^ (mt[state - 1] >> 30)) + state) & MT_WMASK;

        mti_1 = mt[0];
        for(int out = 0, state = 0; out < n_per_RNG; out++) 
        {
            state_1 = state + 1;
            state_M = state + MT_MM;
            if (state_1 >= MT_NN) state_1 -= MT_NN;
            if (state_M >= MT_NN) state_M -= MT_NN;
            mti  = mti_1;
            mti_1 = mt[state_1];
            mti_M = mt[state_M];

            x    = (mti & MT_UMASK) | (mti_1 & MT_LMASK);
            x    =  mti_M ^ (x >> 1) ^ ((x & 1) ? matrix_a : 0);
            mt[state] = x;
            state = state_1;

            //Tempering transformation
            x ^= (x >> MT_SHIFT0);
            x ^= (x << MT_SHIFTB) & mask_b;
            x ^= (x << MT_SHIFTC) & mask_c;
            x ^= (x >> MT_SHIFT1);

            //Convert to (0, 1) float and write to global memory
            // Using unsigned int max, to convert a uniform number in uint range to a uniform range over [-1 ... 1] 
            random[index<2>(out, idx)] = ((float)x + 1.0f) / 4294967296.0f;
        }
    }

    ////////////////////////////////////////////////////////////////////////////////
    // Transform each of MT_RNG_COUNT lanes of n_per_RNG uniformly distributed 
    // random samples, produced by rand_MT_amp(), to normally distributed lanes
    // using Cartesian form of Box-Muller transformation.
    // n_per_RNG must be even.
    ////////////////////////////////////////////////////////////////////////////////
    static void box_muller_transform(float& u1, float& u2) restrict(cpu, amp)
    {
        float r = fast_math::sqrt(-2.0f * fast_math::log(u1));
        float phi = 2.0f * 3.14159265358979f * u2;
        u1 = r * fast_math::cos(phi);
        u2 = r * fast_math::sin(phi);
    }

    static void box_muller_kernel(
        int idx, 
        array_view<float, 2> random, 
        array_view<float, 2> normalized_random, 
        int n_per_RNG
    ) restrict(cpu, amp)
    {
        for(int out = 0; out < n_per_RNG; out += 2) 
        {
            index<2> zero(out, idx);
            index<2> one(out + 1, idx);
            float f0 = random[zero];
            float f1 = random[one];
            box_muller_transform(f0, f1);
            normalized_random[zero] = f0;
            normalized_random[one] = f1;
        }
    }

    void generate_rand_on_amp(
        array_view<float,2> random, 
        bool apply_transform)
    {
        const int n = g_n_per_RNG;
        const array_view<unsigned int, 1>& matrix_a = d_matrix_a;
        const array_view<unsigned int, 1>& mask_b = d_mask_b;
        const array_view<unsigned int, 1>& mask_c = d_mask_c;
        const array_view<unsigned int, 1>& seed = d_seed;

        // generate random numbers
        Concurrency::parallel_for_each(inputsExtent, [=] (index<1> idx) restrict(amp)
        {
            rand_MT_kernel(idx[0], random, matrix_a[idx], mask_b[idx], mask_c[idx], seed[idx], n);
        });

        if (apply_transform)
        {
            std::vector<float> normalized_random_v(randomsExtent.size());
            array_view<float, 2> normalized_random(randomsExtent, normalized_random_v);
            Concurrency::parallel_for_each(inputsExtent, [=] (index<1> idx) restrict(amp)
            {
                box_muller_kernel(idx[0], random, normalized_random, n);
            });
            copy(normalized_random, random); 
        }
    }

    void generate_rand_on_cpu(
        std::vector<float>& v_random, 
        bool apply_transform)
    {
        const int n = g_n_per_RNG;
        const array_view<unsigned int, 1> matrix_a(h_matrix_a.size(), h_matrix_a);
        const array_view<unsigned int, 1> mask_b(h_mask_b.size(), h_mask_b);
        const array_view<unsigned int, 1> mask_c(h_mask_c.size(), h_mask_c);
        const array_view<unsigned int, 1> seed(h_seed.size(), h_seed);
        array_view<float, 2> random(randomsExtent, v_random);

        // generate random numbers
        Concurrency::parallel_for(0, inputsExtent[0], [=] (int idx) restrict(cpu)
        {
            rand_MT_kernel(idx, random, matrix_a[idx], mask_b[idx], mask_c[idx], seed[idx], n);
        });

        if (apply_transform)
        {
            std::vector<float> normalized_random_v(randomsExtent.size());
            array_view<float, 2> normalized_random(randomsExtent, normalized_random_v);
            Concurrency::parallel_for(0, inputsExtent[0], [=] (int idx) restrict(cpu)
            {
                box_muller_kernel(idx, random, normalized_random, n);
            });
            copy(normalized_random, random); 
        }
    }

    MSParallelMersenneTwister()
        :
        h_matrix_a(MT_RNG_COUNT), 
        h_mask_b(MT_RNG_COUNT), 
        h_mask_c(MT_RNG_COUNT), 
        h_seed(MT_RNG_COUNT),
        g_path_n(256),
        //g_n_per_RNG(align_up(div_up(g_path_n, MT_RNG_COUNT), 2)),
        g_n_per_RNG(8),
        g_rand_n( MT_RNG_COUNT * g_n_per_RNG), // how many #'s
        g_seed(777),
        inputsExtent(MT_RNG_COUNT),
        d_matrix_a(inputsExtent),
        d_mask_b(inputsExtent),
        d_mask_c(inputsExtent),
        d_seed(inputsExtent),
        randomsExtent(g_n_per_RNG, MT_RNG_COUNT)
    {
        load_MT(".\\MersenneTwister.dat", data, MT_RNG_COUNT);
        seed_MT(g_seed, data, h_matrix_a, h_mask_b, h_mask_c, h_seed);

        // Copy to GPU
        d_matrix_a = array<unsigned int, 1>(inputsExtent, h_matrix_a.begin());
        d_mask_b = array<unsigned int, 1>(inputsExtent, h_mask_b.begin());
        d_mask_c = array<unsigned int, 1>(inputsExtent, h_mask_c.begin());
        d_seed = array<unsigned int, 1>(inputsExtent, h_seed.begin());
    }

    int total() const { return g_rand_n; }

    mt_struct data[MT_RNG_COUNT];
    std::vector<unsigned int> h_matrix_a, h_mask_b, h_mask_c, h_seed;
    const int g_path_n;
    const int g_n_per_RNG;
    const int g_rand_n; // how many #'s
    const unsigned int g_seed;
    extent<1> inputsExtent;
    array<unsigned int, 1> d_matrix_a, d_mask_b, d_mask_c, d_seed;
    extent<2> randomsExtent;
};
